# -*- coding:utf-8 -*-
# @FileName : app\utils\Constant.py
# @Time     : 2025/11/28
# @Author   : 天空之城
"""全局静态变量工具类"""
from pickle import TRUE
import random
import datetime
import requests
from app.config import MCP_SERVERS_LIST, MCP_SERVERS_CHECK_INTERVAL, logger, mcp_to_cdp

class Constant:
    
    # 任务状态映射字典，key为taskId，value为任务状态
    taskStatusMap = {}
    # 任务MCP映射字典，key为taskId，value为mcp服务地址，也为mcpServerStatusMap.key
    taskMcpServerMap = {}
    # MCP服务器状态映射字典，key为serverUrl，value为服务器状态
    # mcpServerStatusMap = {
    #    "http://127.0.0.1:8931/sse": {"status": "free", "last_time": None, "load": 0},
    #    "http://127.0.0.1:8932/sse": {"status": "free", "last_time": None, "load": 0}
    # }
    mcpServerStatusMap = {server: {"status": "free", "last_time": None, "load": 0} for server in MCP_SERVERS_LIST}
    
    # MCP到CDP的映射关系字典，key为mcp服务地址，value为cdp服务地址
    # 格式：mcp_to_cdp = {'http://127.0.0.1:8931/sse': 'http://127.0.0.1:9222'}
    mcp_to_cdp = mcp_to_cdp.copy()
    
    is_testing_mcp_server = False

    @classmethod
    def set_task_status(cls, task_id, status):
        """设置指定任务ID的状态"""
        cls.taskStatusMap[task_id] = status

    @classmethod
    def get_task_status(cls, task_id):
        """获取指定任务ID的状态"""
        return cls.taskStatusMap.get(task_id)

    @classmethod
    def remove_task_status(cls, task_id):
        """移除指定任务ID的状态记录"""
        if task_id in cls.taskStatusMap:
            del cls.taskStatusMap[task_id]

    @classmethod
    def clear_all_tasks(cls):
        """清除所有任务状态记录"""
        cls.taskStatusMap.clear()

    @classmethod
    def get_all_tasks(cls):
        """获取所有任务状态记录的副本"""
        return cls.taskStatusMap.copy()

    # MCP服务器状态相关方法
    @classmethod
    def set_mcp_server_status(cls, server_url, status):
        """设置指定MCP服务器的状态"""
        if server_url in cls.mcpServerStatusMap:
            cls.mcpServerStatusMap[server_url]["status"] = status

    # 更新、添加MCP服务器
    @classmethod
    def update_mcp_server_status(cls, server_url, server_status):
        """设置指定MCP服务器的状态"""
        cls.mcpServerStatusMap[server_url] = server_status
        logger.info(f"更新MCP服务器 {server_url} 状态为 {server_status['status']}")

    @classmethod
    def get_mcp_server_status(cls, server_url):
        """获取指定MCP服务器的状态"""
        return cls.mcpServerStatusMap.get(server_url)["status"]

    @classmethod
    def set_mcp_server_last_time(cls, server_url, last_time):
        """设置指定MCP服务器的最后使用时间"""
        if server_url in cls.mcpServerStatusMap:
            cls.mcpServerStatusMap[server_url]["last_time"] = last_time

    @classmethod
    def get_mcp_server_last_time(cls, server_url):
        """获取指定MCP服务器的最后使用时间"""
        return cls.mcpServerStatusMap.get(server_url)["last_time"]

    @classmethod
    def remove_mcp_server(cls, server_url):
        """移除指定MCP服务器的状态记录"""
        if server_url in cls.mcpServerStatusMap:
            del cls.mcpServerStatusMap[server_url]

    @classmethod
    def clear_all_mcp_servers(cls):
        """清除所有MCP服务器状态记录"""
        cls.mcpServerStatusMap.clear()

    @classmethod
    def get_all_mcp_servers(cls):
        """获取所有MCP服务器状态记录的副本"""
        return cls.mcpServerStatusMap.copy()
    
    # MCP到CDP映射关系管理方法
    @classmethod
    def set_mcp_to_cdp_mapping(cls, mcp_url, cdp_url):
        """设置MCP到CDP的映射关系
        
        Args:
            mcp_url: MCP服务地址
            cdp_url: CDP服务地址
        """
        cls.mcp_to_cdp[mcp_url] = cdp_url
        logger.info(f"设置MCP到CDP映射关系: {mcp_url} -> {cdp_url}")
    
    @classmethod
    def get_cdp_url_by_mcp(cls, mcp_url):
        """根据MCP服务地址获取对应的CDP服务地址
        
        Args:
            mcp_url: MCP服务地址
            
        Returns:
            str: CDP服务地址，如果没有映射关系则返回None
        """
        return cls.mcp_to_cdp.get(mcp_url)
    
    @classmethod
    def remove_mcp_to_cdp_mapping(cls, mcp_url):
        """移除MCP到CDP的映射关系
        
        Args:
            mcp_url: MCP服务地址
        """
        if mcp_url in cls.mcp_to_cdp:
            del cls.mcp_to_cdp[mcp_url]
            logger.info(f"移除MCP到CDP映射关系: {mcp_url}")
    
    @classmethod
    def get_all_mcp_to_cdp_mappings(cls):
        """获取所有MCP到CDP的映射关系
        
        Returns:
            dict: MCP到CDP的映射关系字典
        """
        return cls.mcp_to_cdp.copy()
    
    # 服务器负载管理方法
    @classmethod
    def increase_server_load(cls, server_url):
        """增加服务器负载
        
        Args:
            server_url: MCP服务地址
        """
        if server_url in cls.mcpServerStatusMap:
            cls.mcpServerStatusMap[server_url]["load"] += 1
            logger.debug(f"增加服务器负载: {server_url}, 当前负载: {cls.mcpServerStatusMap[server_url]['load']}")
    
    @classmethod
    def decrease_server_load(cls, server_url):
        """减少服务器负载
        
        Args:
            server_url: MCP服务地址
        """
        if server_url in cls.mcpServerStatusMap:
            cls.mcpServerStatusMap[server_url]["load"] = max(0, cls.mcpServerStatusMap[server_url]["load"] - 1)
            logger.debug(f"减少服务器负载: {server_url}, 当前负载: {cls.mcpServerStatusMap[server_url]['load']}")
    
    @classmethod
    def get_server_load(cls, server_url):
        """获取服务器负载
        
        Args:
            server_url: MCP服务地址
            
        Returns:
            int: 服务器负载
        """
        return cls.mcpServerStatusMap.get(server_url, {}).get("load", 0)

    @classmethod
    def get_random_free_mcp_server(cls, task_id):
        """根据状态值为free，智能选择一个MCP服务器URL,同时锁定该服务器
        
        智能分配算法：
        1. 优先选择负载较低的服务器
        2. 考虑服务器的可用性和响应时间
        3. 确保任务公平获取资源
        
        Args:
            task_id: 任务ID
            
        Returns:
            str: 可用的MCP服务器URL，如果没有可用服务器则返回None
        """
        now = datetime.datetime.now() 
        
        # 第一次尝试：选择负载最低的可用服务器
        available_servers = []
        for url, status in cls.mcpServerStatusMap.items():
            if status["status"] == 'free':
                available_servers.append((url, status["load"]))
        
        # 按负载排序，优先选择负载较低的服务器
        available_servers.sort(key=lambda x: x[1])
        
        # 尝试可用服务器
        for server_url, _ in available_servers:
            # 先锁定待测试的服务器，防止测试连通性时被其他请求获取
            cls.set_mcp_server_status(server_url, 'lock')
            server_status = cls.mcpServerStatusMap.get(server_url)
            
            # 是否需要连通性测试(守护线程周期内测试过的不需要测试连通性)
            need_test_connect = (not server_status["last_time"]) or \
                               (now - server_status["last_time"]).total_seconds() > MCP_SERVERS_CHECK_INTERVAL
            cls.set_mcp_server_last_time(server_url, now)
            
            # 测试服务器连通性(守护线程周期内测试过的不需要测试连通性)
            is_connected = not need_test_connect or cls._test_server_connectivity(server_url)
            
            if is_connected:
                logger.info(f"{server_url}锁定时间: {now.strftime('%Y-%m-%d %H:%M:%S')}")
                # 标记MCP服务器被占用
                cls.set_mcp_server_status(server_url, 'busy')
                # 增加服务器负载
                cls.increase_server_load(server_url)
                cls.set_task_mcp_server(task_id, server_url)
                return server_url
            else:
                # 标记MCP服务器为fail
                logger.warning(f"{server_url} 连通性测试失败，标记为fail")
                cls.set_mcp_server_status(server_url, 'fail')
        
        # 第二次尝试：尝试状态为fail的服务器（可能已恢复）
        fail_servers = [url for url, status in cls.mcpServerStatusMap.items() if 
                       status["status"] == 'fail' and 
                       (now - status["last_time"]).total_seconds() > 600]  # 10分钟后重试
        
        for fail_server in fail_servers:
            # 先锁定待测试的服务器，防止测试连通性时被其他请求获取
            cls.set_mcp_server_status(fail_server, 'lock')
            cls.set_mcp_server_last_time(fail_server, now)
            
            # 测试服务器连通性
            if cls._test_server_connectivity(fail_server):
                logger.info(f"{fail_server}锁定时间: {now.strftime('%Y-%m-%d %H:%M:%S')}")
                # 标记MCP服务器被占用
                cls.set_mcp_server_status(fail_server, 'busy')
                # 增加服务器负载
                cls.increase_server_load(fail_server)
                cls.set_task_mcp_server(task_id, fail_server)
                return fail_server
            else:
                # 继续标记为fail
                logger.warning(f"{fail_server} 连通性测试失败，继续标记为fail")
                cls.set_mcp_server_status(fail_server, 'fail')
        
        return None
    
    @classmethod
    def _test_server_connectivity(cls, server_url, timeout=2):
        """测试SSE类型接口连通性"""
        try:
            # 为SSE请求设置适当的头信息
            headers = {
                'Accept': 'text/event-stream',
                'Cache-Control': 'no-cache'
            }
            
            # 发送请求但不读取完整响应体，只检查连接是否建立
            with requests.get(server_url, headers=headers, timeout=timeout, stream=True) as response:
                # 检查响应状态码
                if response.status_code != 200:
                    return False
                
                # 检查响应内容类型是否为text/event-stream
                content_type = response.headers.get('Content-Type', '')
                if 'text/event-stream' not in content_type:
                    logger.warning(f"警告: {server_url} 不是SSE接口，Content-Type: {content_type}")
                    return True  # 即使不是标准SSE类型，只要能连接上也返回True
                
                # 尝试读取一点响应内容，确认连接正常
                try:
                    # 读取一行数据但不阻塞太长时间
                    for line in response.iter_lines(decode_unicode=True, chunk_size=1):
                        if line:
                            return True
                        # 如果没有数据立即返回，也认为连接成功
                        return True
                except:
                    # 如果读取过程中出现异常，也认为连接成功（只要能建立连接）
                    return True
        except Exception as e:
            # 连接失败
            logger.error(f"SSE连接测试异常: {e}")
            return False

    @classmethod
    def set_task_mcp_server(cls, task_id, mcp_server_url):
        """设置指定任务ID的MCP服务器"""
        cls.taskMcpServerMap[task_id] = mcp_server_url

    @classmethod
    def get_task_mcp_server(cls, task_id):
        """获取指定指定任务ID的MCP服务器"""
        return cls.taskMcpServerMap.get(task_id)

    @classmethod
    def remove_task_mcp_server(cls, task_id):
        """移除指定MCP服务器的状态记录"""
        if task_id in cls.taskMcpServerMap:
            del cls.taskMcpServerMap[task_id]

    @classmethod
    def release_task_mcp_server(cls, task_id):
        """释放指定任务ID的MCP服务器
        
        Args:
            task_id: 任务ID
        """
        task_server = cls.get_task_mcp_server(task_id)
        if task_server:
            cls.set_mcp_server_status(task_server, 'free')
            # 减少服务器负载
            cls.decrease_server_load(task_server)
            now = datetime.datetime.now() 
            logger.info(f"{task_server}释放时间: {now.strftime('%Y-%m-%d %H:%M:%S')}")
            cls.set_mcp_server_last_time(task_server, now)
            cls.remove_task_mcp_server(task_id)
    
    @classmethod
    def release_task_mcp_server_by_url(cls, task_server):
        """释放指定task_server的MCP服务器
        
        Args:
            task_server: MCP服务器URL
        """
        if task_server:
            cls.set_mcp_server_status(task_server, 'free')
            # 减少服务器负载
            cls.decrease_server_load(task_server)
            now = datetime.datetime.now() 
            logger.info(f"{task_server}释放时间: {now.strftime('%Y-%m-%d %H:%M:%S')}")
            cls.set_mcp_server_last_time(task_server, now)
            # 根据task_server从taskMcpServerMap移除对应任务
            for task_id, server_url in list(cls.taskMcpServerMap.items()):
                if server_url == task_server:
                    cls.remove_task_mcp_server(task_id)
                    break

    @classmethod
    def get_free_mcp_server_count(cls):
        """获取状态为free的MCP服务器数量"""
        free_count = 0
        for server_status in cls.mcpServerStatusMap.values():
            if server_status["status"] == "free":
                free_count += 1
        return free_count
    
    @classmethod
    def check_all_mcp_servers_availability(cls):
        """全局MCP工具可用性检查
        
        检查所有MCP服务器的可用性，更新服务器状态，确保用户每次发起任务前
        都能获取到可用的MCP执行器。
        
        Returns:
            dict: 包含检查结果的字典，格式为{"available_servers": [...], "total_servers": int, "available_count": int}
        """
        now = datetime.datetime.now()
        available_servers = []
        total_servers = len(cls.mcpServerStatusMap)
        
        logger.info(f"开始全局MCP服务器可用性检查，共{total_servers}个服务器")
        
        for server_url, server_status in cls.mcpServerStatusMap.items():
            # 跳过正在使用的服务器
            if server_status["status"] == "busy":
                logger.debug(f"跳过正在使用的服务器: {server_url}")
                continue
            
            # 检查服务器连通性
            is_connected = cls._test_server_connectivity(server_url)
            
            if is_connected:
                # 更新服务器状态为可用
                if server_status["status"] != "free":
                    cls.set_mcp_server_status(server_url, "free")
                    logger.info(f"服务器恢复可用: {server_url}")
                available_servers.append(server_url)
            else:
                # 更新服务器状态为不可用
                if server_status["status"] != "fail":
                    cls.set_mcp_server_status(server_url, "fail")
                    logger.warning(f"服务器不可用: {server_url}")
            
            # 更新最后检查时间
            cls.set_mcp_server_last_time(server_url, now)
        
        available_count = len(available_servers)
        logger.info(f"全局MCP服务器可用性检查完成，可用服务器: {available_count}/{total_servers}")
        
        return {
            "available_servers": available_servers,
            "total_servers": total_servers,
            "available_count": available_count
        }
    
    @classmethod
    def ensure_available_mcp_server(cls, task_id):
        """确保有可用的MCP服务器，若没有则执行全局检查
        
        Args:
            task_id: 任务ID
            
        Returns:
            str: 可用的MCP服务器URL，如果没有可用服务器则返回None
        """
        # 先尝试获取可用服务器
        server_url = cls.get_random_free_mcp_server(task_id)
        
        # 如果没有可用服务器，执行全局可用性检查后再次尝试
        if not server_url:
            logger.warning("没有可用的MCP服务器，执行全局可用性检查")
            cls.check_all_mcp_servers_availability()
            server_url = cls.get_random_free_mcp_server(task_id)
        
        return server_url

